#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include "lab5.h"

#define PACKETSIZE (256)
//#define DEBUG (1)
#define ALPHA (1/8)
#define BETA (1/4)

// Connection state variables
int sequence_number;
int senderLastAck;
int senderCounter;
int currentRTT;
int currentVar;
int currentTimeout;
int receiverNextAck;
int finalAckNum;
int receivedFin;
int initiatedClose;

int windowSize;


/* The following three helper functions make dealing with time values a bit
 * easier.  You probably don't need to modify them, but you can if you want. */
int timeval_to_msec(struct timeval *t) { 
    return t->tv_sec*1000+t->tv_usec/1000;
}

void msec_to_timeval(int millis, struct timeval *out_timeval) {
    out_timeval->tv_sec = millis/1000;
    out_timeval->tv_usec = (millis%1000)*1000;
}

int current_msec() {
    struct timeval t;
    gettimeofday(&t,0);
    return timeval_to_msec(&t);
}
/* End time functions. */


int my_connect(int socket,struct sockaddr_in *toaddr,int addrsize) {
    return connect(socket,(struct sockaddr*)toaddr,addrsize);
}

int my_rtt(int socket) {
  return currentRTT;
}

void update_rtt( int rtt, int timeout ) {
  if ( timeout ) {
    currentTimeout = 2 * currentTimeout;
  }
  else {
    double cRTT = (0.875)*currentRTT + (0.125)*rtt;
    double cVar = (0.75) * currentVar;
    if (rtt > currentRTT) {
      cVar += 0.25 * (rtt - currentRTT);
    }
    else {
      cVar += 0.25 * (currentRTT - rtt);
    }
    currentRTT = (int) cRTT;
    currentVar = (int) cVar;
    currentTimeout = currentRTT + 4 * currentVar;
  }
#ifdef DEBUG
  fprintf(stderr, "\nUpdateRTT: Samp-%d  Timeout-%d, NewTO-%d, newRTT-%d, newVar- %d", rtt, timeout, currentTimeout, currentRTT, currentVar);
  #endif
}


void my_send( int sock, void * buf, size_t len ) {
  my_sendf( sock, buf, len, 0, -1 );
}

void my_sendf(int sock, void *buf, size_t len, int fin, int maxTries)
{

  // We want to package up the data with a header that includes a timestamp,
  // a packet number, and a sequence number. 
  
  // Setup our sets for select
  fd_set rfds, wfds;
  FD_ZERO(&rfds);
  FD_ZERO(&wfds);
  FD_SET(sock, &rfds);

  int numRemaining = len;
  int retransmitted = 0;
  finalAckNum += len;

  char * packet[MAX_PACKET];
  char * receive[MAX_PACKET];
  struct lab5_hdr * hdr;
  int length; 
  while ( senderLastAck < finalAckNum ) {

    // If this is the last time, we might not get an ack back
    if ( maxTries > 0 && retransmitted >= maxTries ) {
      // Try to send the final fin maxTries times
      return;
    }

    FD_ZERO(&rfds);
    FD_ZERO(&wfds);
    FD_SET(sock, &rfds);
    
    // We only need to rebuild the packet if we are not retransmitting
    // an existing one
    if ( ! retransmitted ) {
      /* Buffer for the packet. */
      memset(packet, 0, MAX_PACKET * sizeof(char));
      
      
      hdr = (struct lab5_hdr *) packet;
      
      if ( MAX_PACKET - sizeof(struct lab5_hdr) > numRemaining ) {
	length = numRemaining;
      }
      else {
	length = MAX_PACKET - sizeof(struct lab5_hdr);
      }
      
      /* Set the header fields as needed. */
      hdr->packet_number = htons(senderCounter);
      hdr->length = htons(length);
      hdr->timestamp = htonl(current_msec());
      hdr->sequence_number = htonl(sequence_number);
      hdr->ack_flag = htons( 0 );
      if ( fin ) {
	hdr->fin_flag = htons( 1 );
      }


      /* Copy the payload into our packet.  hdr + 1 is the address of the
       * first byte after the header. */
      memcpy(hdr+1, buf + len - numRemaining,length);
    }
    else {
      // Do change the packet number (for debugging purposes)
      hdr->packet_number = htons(senderCounter);;
    }
#ifdef DEBUG
      fprintf(stderr, "\nSending Packet: Num-%d  Len-%d  Time-%d  Seq-%d  Ack-0 Fin-%d",
	      senderCounter, length, ntohl(hdr->timestamp), sequence_number, fin);
#endif


    // Construct a timeval spec
    struct timeval t;
    msec_to_timeval( currentTimeout, &t );
    
    send(sock, packet, sizeof(struct lab5_hdr) + length, 0);
    senderCounter ++;  

#ifdef DEBUG
    fprintf(stderr, "\nWaiting on Select()");
#endif

    // Now, we need to call select() to wait to receive an ack
    int ret = select( sock + 1, &rfds, &wfds, NULL, &t );
    if (ret < 0 ) {
      perror("select");
    }
    if (ret == 0) { // Timeout
#ifdef DEBUG
      fprintf(stderr, "\nSelect() timeout waiting for ack. Timeout %d", currentTimeout);
#endif
      retransmitted += 1;
      update_rtt( 0, 1 );
    }
    else {
#ifdef DEBUG
      fprintf(stderr, "\nSelect() - Reading in data");
#endif
      // Try to interpret the ack, discarding if anything seems invalid
      struct sockaddr_in fromaddr;
      socklen_t addrlen = sizeof(fromaddr);
      memset(receive, 0, MAX_PACKET * sizeof(char));
      int recv_count = recvfrom(sock, receive, MAX_PACKET, 0,
				(struct sockaddr *) &fromaddr, &addrlen);
      if ( recv_count < sizeof(struct lab5_hdr) ) {
	continue;
      }
#ifdef DEBUG
      fprintf(stderr, "\nReceived %d bytes. (%lu acks)", recv_count, recv_count / sizeof(struct lab5_hdr));
#endif


      struct lab5_hdr *ack = (struct lab5_hdr *) receive;
      int packetNum = ntohs(ack->packet_number);
      int timeDiff = current_msec() - ntohl(ack->timestamp);
      int nextExpect = ntohl(ack->sequence_number);
      int ackn = ntohs(ack->ack_flag);
      int finf = ntohs(ack->fin_flag);
#ifdef DEBUG
      fprintf(stderr, "\nReceived Packet: Num-%d  RTT-%d  Syn-%d  Ack-%d  Fin-%d",
	      packetNum, timeDiff, nextExpect, ackn, finf);
#endif
      // If we get a fin while trying to send, then a prior ack must have 
      // been lost. So, go ahead and return, since they must have gotten
      // our fin
      if ( finf && initiatedClose ) {
	receivedFin = 1;
	return;
      }
	
      if ( ackn > 0 ) {
	senderLastAck = nextExpect;
	sequence_number = nextExpect;
	// We need to send until we reach finalAckNum
	numRemaining = finalAckNum - sequence_number;
	if ( ! retransmitted ) {
	  update_rtt( timeDiff, 0 );
	}
	// Now we're sending a new packet
	retransmitted = 0;

	

      } /* End of if block for being an ack */
    } /* End of else block for not timing out */
  } /* End of while loop for unsent data */
}

int my_socket(int domain, int type, int protocol) {
    /* Initialize state. */
    sequence_number = 0;
    currentRTT = 1000;
    currentTimeout = 1000;
    currentVar = 100;
    receiverNextAck = 0;
    finalAckNum = 0;
    receivedFin = 0;
    initiatedClose = 0;
    return socket(domain, type, protocol);
}


int my_recv(int sock, void * buffer, size_t length) {

  // We need to
  // 1) Extract information from the packet, including the sequence
  //    number and the length.
  // 3) See if our sequence number matches what we're expecting
  //    -- If it does,    then return the data and ack the data
  //       keeping track of the next sequence number that we're expecting
  //    -- If it doesn't, then do nothing 

    /* Buffer for the packet. */
    char packet[MAX_PACKET];

    while ( 1 ) {
      memset(packet, 0, sizeof(packet));

      struct lab5_hdr *hdr = (struct lab5_hdr *) packet;
      
      struct sockaddr_in fromaddr;
      socklen_t addrlen = sizeof(fromaddr);
      int recv_count = recvfrom(sock, packet, MAX_PACKET, 0,
				(struct sockaddr *) &fromaddr, &addrlen);
      
      // TODO - handle case of a short receive?
      // Maybe read in the buffer size and then try to read in the specified
      // length. If it's short, then throw out the packet. 
      
      /* Associate this UDP socket with a remote address.  This is simply
       * for convenience so that we can use send() rather than sendto(), since
       * sendto() requires an additional argument that we don't want to deal 
       * with every time we need to send something. */
      if(connect(sock, (struct sockaddr *) &fromaddr, addrlen)) {
        perror("connect (my_recv)");
      }
#ifdef DEBUG
      fprintf(stderr, "\nReceived: Num-%d  Len-%d  Seq-%d  Expect-%d", 
	      ntohs(hdr->packet_number), ntohs(hdr->length),
	      ntohl(hdr->sequence_number), receiverNextAck);
#endif
      if ( initiatedClose ) {
	// Discard everything until we get a fin response
	if ( ntohs(hdr->fin_flag) ) {
	  receiverNextAck += recv_count - sizeof(struct lab5_hdr);
	  sendAck ( sock, receiverNextAck, packet );
	  return 0;
	}
	continue;
      }

      // Throw out the packet if it is not the expected length
      if ( ntohs(hdr->length) != recv_count - sizeof(struct lab5_hdr) ) {
#ifdef DEBUG
	fprintf(stderr, "\nReceived unexpected length packet. Hdr-%d  Got-%lu", ntohs(hdr->length), recv_count - sizeof(struct lab5_hdr));
#endif
	continue;
      }
      
#ifdef DEBUG
	fprintf(stderr, "\nPacket received has expected length: Hdr-%d  Got-%lu", ntohs(hdr->length), recv_count - sizeof(struct lab5_hdr));
#endif


      // Extract the sequence number from the packet, since we want to
      // see if it's the number we're expecting
      if ( ntohl(hdr->sequence_number) == receiverNextAck ) {
#ifdef DEBUG
	fprintf(stderr, "\nPacket received has expected sequence number.");
#endif
	// This is a good packet
	
	// Only copy over the data if we are not closing 
	if ( 0 == ntohs(hdr->fin_flag) ) {
	  /* Copy the payload into the user-supplied buffer. */
	  memcpy(buffer, packet + sizeof(struct lab5_hdr),
		 recv_count - sizeof(struct lab5_hdr));
	}
	receiverNextAck += recv_count - sizeof(struct lab5_hdr);
	sendAck ( sock, receiverNextAck, packet );
	if ( 0 == ntohs(hdr->fin_flag) ) {
	  return recv_count - sizeof(struct lab5_hdr);
	}
	else {
	  receivedFin = 1;
	  return 0;
	}
	
      }
#ifdef DEBUG
      fprintf(stderr, "\nPacket received has unexpected sequence number.");
#endif
      // Sender didn't get our last ack, so resend it
      sendAck( sock, receiverNextAck, packet );
      // And loop around again to receive the next packet
    } 
    // Never get here
    return -1;
}


void sendAck( int sock, int ackNum, char * packet ) {
  char ack[sizeof(struct lab5_hdr)];
  struct lab5_hdr * ackHeader = (struct lab5_hdr *) ack;
  // We only need to change a few values here
  memcpy( ack, packet, sizeof(struct lab5_hdr) );
  ackHeader->length = htons(0);
  ackHeader->sequence_number = htonl( ackNum );
  ackHeader->ack_flag = htons(1);

  // Don't piggyback fins and acks
  ackHeader->fin_flag = htons(0);

  // Send the Ack
  send(sock, ackHeader, sizeof(struct lab5_hdr), 0);
#ifdef DEBUG
  fprintf(stderr, "\nSent Ack for Seq Number - %d", ackNum );
#endif
}

int my_close(int sock) {
    /* Send an empty packet to signal end of file. */
  char buf[128];
  memset(buf, 0, 128*sizeof(char));

  if ( 0 == receivedFin ) {
#ifdef DEBUG
    fprintf(stderr, "\n** Closing Initiator **");
#endif
    initiatedClose = 1;
    // Send a fin packet and get an ack for it
    my_sendf( sock, buf, 128, 1, -1 );
#ifdef DEBUG
    fprintf(stderr, "\n** Closing Initiator :: Sent Fin and Got Ack **");
#endif
    my_recv( sock, buf, 128 );
#ifdef DEBUG
    fprintf(stderr, "\n** Closing Initiator :: Received Fin **");
#endif
  }
  else {
#ifdef DEBUG
    fprintf(stderr, "\n** Closing Responder **");
#endif
    // We need to send a fin packet at most 3 times
    my_sendf( sock, buf, 128, 1, 3 );
#ifdef DEBUG
    fprintf(stderr, "\n** Closing Responder :: Sent Fin Response **");
#endif
  }
  return close(sock);

  // If we are initiating the close then
  // Send a fin packet and wait for an ack
  // Then, wait until we receive a fin packet
  // Ack that packet
  // Close the socket
  // Return

  // If we have received a fin packet (we are not initiaing) then
  // Ack any fin packets that come in
  // Send a fin packet up to 3 times, waiting for an ack or returning 0
  // after the 3rd time. 


  //my_sendf(sock, "CLOSE", strlen("CLOSE"), 1);
  // Send the last ack and close the connection
  //char packet[sizeof(struct lab5_hdr)];
  //sendAck( sock, 0, packet );
  //sleep(5);
  
}
